import cv2
# import numpy as np
from ultralytics import YOLO
import time
# import matplotlib.pyplot as plt

def main():
    # 初始化YOLOv8模型（自动下载预训练权重）
    model = YOLO('yolov8n.pt')  # 使用轻量级版本，适合实时检测

    # 打开视频源 - 可以选择摄像头、视频文件或网络流
    # video_source = 0  # 默认摄像头
    video_source = "traffic.mp4"  # 视频文件
    cap = cv2.VideoCapture(video_source)

    if not cap.isOpened():
        print("无法打开视频源")
        return

    # 获取视频基本信息
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    print(f"视频尺寸: {width}x{height}, FPS: {fps}")

    # 创建视频保存对象（可选）
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    out = cv2.VideoWriter('output.avi', fourcc, 20.0, (width, height))
    print("视频保存路径: output.avi")

    # 车辆类别ID（COCO数据集中）
    vehicle_classes = [2, 3, 5, 7]  # car, motorcycle, bus, truck

    # 性能统计
    frame_count = 0
    start_time = time.time()

    while True:
        print("正在处理帧...")
        ret, frame = cap.read()
        # print(ret, frame)
        if not ret:
            break

        frame_count += 1

        # 使用YOLOv8进行目标检测
        results = model.track(frame, persist=True, verbose=False, classes=vehicle_classes)

        # 获取检测结果
        boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
        classes = results[0].boxes.cls.cpu().numpy().astype(int)
        confidences = results[0].boxes.conf.cpu().numpy()
        track_ids = results[0].boxes.id.cpu().numpy().astype(int) if results[0].boxes.id is not None else None

        # 绘制检测结果
        for i, box in enumerate(boxes):
            x1, y1, x2, y2 = box

            # 获取类别名称和置信度
            class_id = classes[i]
            confidence = confidences[i]

            # 设置颜色和标签
            color = get_class_color(class_id)
            label = f"{model.names[class_id]} {confidence:.2f}"

            # 绘制边界框
            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 2)

            # 绘制标签背景
            (label_width, label_height), baseline = cv2.getTextSize(
                label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1
            )
            cv2.rectangle(
                frame,
                (x1, y1 - label_height - 5),
                (x1 + label_width, y1 - 2),
                color,
                -1
            )

            # 绘制标签文本
            cv2.putText(
                frame,
                label,
                (x1, y1 - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                (255, 255, 255),
                1
            )

            # 绘制跟踪ID（如果可用）
            if track_ids is not None:
                track_id = track_ids[i]
                cv2.putText(
                    frame,
                    f"ID: {track_id}",
                    (x1, y1 - 30),
                    cv2.FONT_HERSHEY_SIMPLEX,
                    0.6,
                    (0, 255, 255),
                    2
                )

        # 显示帧率
        elapsed_time = time.time() - start_time
        fps = frame_count / elapsed_time
        cv2.putText(
            frame,
            f"FPS: {fps:.2f}",
            (10, 30),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 0, 255),
            2
        )

        # 显示车辆计数
        vehicle_count = len(boxes)
        cv2.putText(
            frame,
            f"Vehicles: {vehicle_count}",
            (10, 70),
            cv2.FONT_HERSHEY_SIMPLEX,
            1,
            (0, 255, 0),
            2
        )

        # 显示处理后的帧
        cv2.imshow("Vehicle Detection", frame)

        # plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        # plt.title("Vehicle Detection")
        # plt.pause(0.001)  # 短暂暂停以更新显示
        # plt.clf()  # 清除当前图形


        # 保存视频（可选）
        out.write(frame)
        print("视频保存")

        # 按'q'退出
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放资源
    cap.release()
    # out.release()
    cv2.destroyAllWindows()

    print(f"处理完成！总帧数: {frame_count}, 平均FPS: {fps:.2f}")


def get_class_color(class_id):
    """为不同车辆类型分配不同颜色"""
    colors = {
        2: (0, 255, 0),  # 小汽车 - 绿色
        3: (255, 0, 0),  # 摩托车 - 蓝色
        5: (0, 165, 255),  # 公交车 - 橙色
        7: (0, 0, 255)  # 卡车 - 红色
    }
    return colors.get(class_id, (255, 255, 0))  # 默认为黄色


if __name__ == "__main__":
    main()